"""
预处理程序
在爬虫开始前进行预处理
1. 建立SQLite数据库，存储爬取状态
2. 建立爬虫结果文件，并写入表头
3. 建立日志文件
"""

import csv
import sqlite3
import pandas as pd

from config import csv_path, sql_path, table_name, tmp_path, log_path, heads


def load_sql(csv_path, sql_path, table_name):
    """将企业名称写进SQLite数据库，用于断点重爬

    :param csv_path: csv源文件路径
    :param sql_path: SQLite文件路径
    :param table_name: 表名称
    :return:
    """
    with sqlite3.connect(sql_path) as conn:
        # 读取企业名称
        df = pd.read_csv(csv_path)

        # 初始化状态
        df['done'] = 0

        # 写入数据库
        df.to_sql(table_name, conn, if_exists='append', index=False)


def create_tmp(tmp_path, heads):
    """给爬虫结果文件写入表头

    :param tmp_path: 爬虫结果路径
    :param heads: 表头
    :return:
    """
    with open(tmp_path, 'w', encoding="utf8") as f:
        f.write(",".join(heads)+"\n")


if __name__ == '__main__':
    z = input("是否要初始化？已爬取的任何结果将被清空！！！\n输入’Yes‘确认初始化：")
    if z == 'Yes':
        load_sql(csv_path, sql_path, table_name)
        create_tmp(tmp_path, heads)
        with open(log_path, 'w', encoding='utf8') as f:
            pass
    else:
        print("退出初始化")
